BERTの日本語事前学習済みモデルでテキスト埋め込みをやってみる

osawa yuto

2019.04.04

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

どうも、大阪DI部の大澤です。

汎用言語表現モデルBERTの日本語Wikipediaで事前学習済みのモデルがあったので、BERTモデルを使ったテキストの埋め込みをやってみたいと思います。

BERT日本語Pretrainedモデル - KUROHASHI-KAWAHARA LAB

以下のエントリではbert-as-serviceを使ったテキストの埋め込みを紹介しました。今回はBERTのリポジトリで公開されているextract_features.pyを使って、テキストの埋め込みを試します。

bert-as-serviceを使って日本語の文章を埋め込み表現ベクトルに変換してみる｜ DevelopersIO

BERT(Bidirectional Encoder Representations from Transformers)

BERTはGoogleが開発した汎用言語表現モデルです。機械学習で言語表現を学習させることができるため、言語に関わる様々な問題に使う事ができます。

[1810.04805] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

モデル構成の概要

BERTはTransformerと言う機械翻訳モデルのエンコーダー部分を重ねたものになります。ベースモデルであれば12層なので、Transformerが12個重なります。Transformerの前段でトークンと文章とトークンの位置それぞれを埋め込み、その埋め込み表現がTransformerに入力されて、最終的にトークンそれぞれの言語表現(埋め込み表現)が出てくる感じです。

入力されるトークンの先頭は[CLS]という特殊なトークンで、このトークンに対応する言語表現が文章全体の言語表現に相当します。なので、今回の目的である文章の埋め込み表現を取得するときは、このトークン部分を参照することになります。

学習済みモデルの概要

今回使用するのは京都大学の黒橋・河原研究室で公開されている学習済みBERTモデルです。モデルの紹介ページでは次のように説明されています。

入力テキスト: 日本語Wikipedia全て (約1,800万文)

入力テキストにJuman++ (v2.0.0-rc2)で形態素解析を行い、さらにBPEを適用しsubwordに分割

BERT_{BASE}と同じ設定 (12-layer, 768-hidden, 12-heads)

30 epoch (1GPUで1epochに約1日かかるのでpretrainingに約30日)

語彙数: 32,000 (形態素、subwordを含む)

max_seq_length: 128

事前学習(pretraining)に約30日...! こういったリソースを公開して頂けるのは非常に有難いですね。

やってみる

環境

macOS Mojave
Python 3.7.2

※ TensorFlowなどPythonのライブラリが必要になります。進めていく中でエラーが出る場合は必要に応じて、ライブラリをインストールしてください。

スクリプトの準備

まずは、以下のページから学習済みモデルをダウンロードします。zip形式なので、解凍して適当な場所に配置します。

BERT日本語Pretrainedモデル - KUROHASHI-KAWAHARA LAB

BERTのリポジトリをクローンしてきます。

git clone https://github.com/google-research/bert

今回はJUMAN++(pyknp)でトークン化するため、bert/tokenization.pyを少し修正します。

JUMAN++を使ってトークン化するためのクラス定義JumanPPTokenizerをスクリプトの最後に追加します。

class JumanPPTokenizer(BasicTokenizer):
  def __init__(self):
    """Constructs a BasicTokenizer.
    """
    from pyknp import Juman

    self.do_lower_case = False
    self._jumanpp = Juman()

  def tokenize(self, text):
    """Tokenizes a piece of text."""
    text = convert_to_unicode(text.replace(' ', ''))
    text = self._clean_text(text)

    juman_result = self._jumanpp.analysis(text)
    split_tokens = []
    for mrph in juman_result.mrph_list():
      split_tokens.extend(self._run_split_on_punc(mrph.midasi))

    output_tokens = whitespace_tokenize(" ".join(split_tokens))
    print(split_tokens)
    return output_tokens

tokenization.py内のFullTokenizerを修正します。元々BasicTokenizerを使っていた箇所をJumanPPTokenizerに変更します。

  def __init__(self, vocab_file, do_lower_case=True):
    self.vocab = load_vocab(vocab_file)
    self.inv_vocab = {v: k for k, v in self.vocab.items()}
    # self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
    self.jumanpp_tokenizer = JumanPPTokenizer()
    self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)

  def tokenize(self, text):
    split_tokens = []
    # for token in self.basic_tokenizer.tokenize(text):
    for token in self.jumanpp_tokenizer.tokenize(text):
      for sub_token in self.wordpiece_tokenizer.tokenize(token):
        split_tokens.append(sub_token)

    return split_tokens

extract_features.pyで出力されるテキストはjsonl形式なので、そこから埋め込み表現ベクトルを取り出して、tsvに書き出すスクリプトを用意しておきます。

import json
import numpy as np

# 参照するレイヤーを指定する
TARGET_LAYER = -2

# 参照するトークンを指定する
SENTENCE_EMBEDDING_TOKEN = '[CLS]'

with open('/tmp/output.jsonl', 'r') as f:
    output_jsons = f.readlines()

embedding_list = []
for output_json in output_jsons:
    output = json.loads(output_json)
    for feature in output['features']:
        if feature['token'] != SENTENCE_EMBEDDING_TOKEN: continue
        for layer in feature['layers']:
            if layer['index'] != TARGET_LAYER: continue
            embedding_list.append(layer['values'])

np.savetxt('/tmp/output.tsv', embedding_list, delimiter='\t')

文章を埋め込む

準備ができたので、文章を埋め込んでみます。

extract_features.pyの引数は次の通りです。

input_file: 入力するテキストを改行区切りで記述したテキストファイル
output_file: 出力結果を書き出すjsonl形式のファイル
vocab_file: 学習した単語(トークン)を記述したファイル(トークン化の際に使用)
- 学習済みモデルに含まれる
bert_config_file: BERTモデルの設定ファイル
- 学習済みモデルに含まれる
init_checkpoint: BERTモデルに読み込むチェックポイントファイル(学習した重み付けなど)
- 学習済みモデルに含まれる
do_lower_case: トークン化時に文字を小文字にするかどうか
layers: 必要なTransformerのエンコーダーは何層目のものか
- 複数のレイヤーの結果が欲しい場合はカンマ区切りで指定すればOK(例:-1,-2,-3)
- 基本的には後の方のレイヤーを使うことが多いと思うのでマイナスの値が便利(-1が最後のレイヤー)
- 最後のレイヤーだと事前学習のタスクの影響が大きすぎるかと思い、今回は最後から2つ目のレイヤーの出力を埋め込みに使用しています(とはいえ、可視化した図を見るとそこまで違いはなかったです)

今回使用するもの以外にもいくつか引数があります。詳細はスクリプトbert/extract_features.pyをご覧ください。

cat << EOS > /tmp/input.txt
こんにちは
おはよう
こんばんは
お腹すいた
ご飯食べたい
明日晴れてると良いな
明日の天気はどうだろうか
雨降ったら嫌やな
EOS

export BERT_BASE_DIR="/Users/osawa.yuto/git/ml_pr/sandbox/bert/models/Japanese_L-12_H-768_A-12_E-30_BPE"

python ./bert/extract_features.py \
  --input_file=/tmp/input.txt \
  --output_file=/tmp/output.jsonl \
  --vocab_file=$BERT_BASE_DIR/vocab.txt \
  --bert_config_file=$BERT_BASE_DIR/bert_config.json \
  --init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
  --do_lower_case False \
  --layers -2

数秒から数十秒くらいで処理が完了します。コンソールに出力される情報はかなり多いですが、エラーが吐かれて無ければ問題ないかと思います。

出力されたファイルを確認してみます。形式は次のようになっています。jsonlなので、文章ごとに以下のjsonが一行ずつあります。

{
    "linex_index": 入力された文章が何行目のものか(0ベース),
    "features": [
        {
            "token": 対応するトークン,
            "layers": [
                {
                    "index": レイヤーの番号,
                    "values": 対応する表現ベクトル
                }
            ]
        }
    ]
}

head -1 /tmp/output.jsonl

jsonlから表現ベクトルだけ取り出してtsv形式で書き出します。

python convert_embedding_jsonl.py
head -1 /tmp/output.tsv

可視化

Embedding Projectorで各文章のマッピングを可視化してみたいと思います。先ほど作成した/tmp/input.txtと/tmp/output.tsvを入力(Load)することで、簡単に可視化できます。今回はT-SNEで次元圧縮して3次元にマッピングしました。

Embedding projector - visualization of high-dimensional data

塊ごとにみていきます。天気に関連した文章が近くに集まっています。

"お腹すいた"と"ご飯食べたい"も近くにマッピングされていますね。

挨拶系の言葉も近くにまとまっています。

全体的に直感的なイメージに近いマッピングになっていました。

さいごに

今回は日本語Wikipediaで学習済みのBERTモデルを使用して、文章の埋め込みを行いました。もともと公式リポジトリでスクリプトが用意されていたため、トークン化用のスクリプトを修正するだけで出来ました。BERTは色々な自然言語タスクに応用することができるようなので、今後も何か試してみたいと思います。